package webtrans

import (
	"bytes"
	"gitee.com/DreamPublic/signal/pkg/h3"
	"gitee.com/DreamPublic/signal/pkg/proto"
	"github.com/lucas-clemente/quic-go"
	"github.com/lucas-clemente/quic-go/quicvarint"
	"go.uber.org/zap"
	"io"
	"sync"
)

type Session struct {
	id     string
	cbLock sync.RWMutex
	// header - []proto.SessionCallback
	onCbMap     sync.Map
	quicSession quic.Session
	bidStream   quic.Stream
	signal      *Signal
	log         *zap.Logger
}

func (s *Session) ID() string {
	return s.id
}

func (s *Session) On(header string, cb proto.SessionCallback) {
	s.cbLock.Lock()
	defer s.cbLock.Unlock()
	cbsRaw, ok := s.onCbMap.Load(header)
	if !ok {
		cbsRaw = make([]proto.SessionCallback, 0)
	}
	cbs := append(cbsRaw.([]proto.SessionCallback), cb)
	s.onCbMap.Store(header, cbs)
}

func (s *Session) Emit(header string, body string) error {
	msgStruct := proto.NewMsgStructFromParam(header, body)
	_, writeErr := s.bidStream.Write(msgStruct.ToBytes())
	if nil != writeErr {
		s.log.Error("发送返回数据失败", zap.Error(writeErr),
			zap.String("header", header),
			zap.String("body", body))
	}
	return writeErr
}

func (s *Session) Close() {
	_ = s.quicSession.CloseWithError(0, "")
	s.signal.fireCloseCb(s)
	s.log.Info("连接断开")
}

func NewSession(id string, quicSession quic.Session, signal *Signal) *Session {
	return &Session{
		id:          id,
		cbLock:      sync.RWMutex{},
		onCbMap:     sync.Map{},
		quicSession: quicSession,
		signal:      signal,
		log:         zap.L().With(zap.String("sessionId", id)),
	}
}

// HandleWebTrans 处理webTransport相关逻辑
func (s *Session) HandleWebTrans() {
	defer s.Close()
	// 循环处理连接,异常结束时,关闭连接
	log := s.log
	ctx := s.quicSession.Context()
	// 服务器先发SettingFrame
	uniStreamFromServer, openUniStreamFromServerErr := s.quicSession.OpenUniStreamSync(ctx)
	if nil != openUniStreamFromServerErr {
		log.Error("open UniStreamFromServer error", zap.Error(openUniStreamFromServerErr))
		return
	}
	// 发送流类型
	buf := bytes.Buffer{}
	quicvarint.Write(&buf, h3.ControlStreamType)
	_, writeStreamTypeErr := uniStreamFromServer.Write(buf.Bytes())
	if nil != writeStreamTypeErr {
		log.Error("write StreamType error ", zap.Error(writeStreamTypeErr))
		return
	}
	// 发送SettingFrame
	settingFrameFromServer := h3.SettingFrame{Payload: map[uint64]uint64{
		h3.SettingH3Datagram:         uint64(1),
		h3.SettingEnableWebTransport: uint64(1),
	}}
	_, writeSettingFrameErr := uniStreamFromServer.Write(settingFrameFromServer.ToBytes())
	if nil != writeSettingFrameErr {
		log.Error("write SettingFrame error", zap.Error(writeSettingFrameErr))
		return
	}
	uniStreamFromClient, acceptUniStreamFromClient := s.quicSession.AcceptUniStream(ctx)
	if nil != acceptUniStreamFromClient {
		log.Error("accept UniStreamFromClient error", zap.Error(acceptUniStreamFromClient))
		return
	}
	// 读取流类型
	reader := quicvarint.NewReader(uniStreamFromClient)
	streamType, _ := quicvarint.Read(reader)
	if h3.ControlStreamType == streamType {
		log.Info("收到客户端的ControlStream")
	} else {
		log.Error("客户端发送的uniStream不是ControlStream,关闭")
		return
	}

	// 读取客户端的SettingFrame
	settingFrameFromClient := h3.ParseSettingFrame(uniStreamFromClient)
	log.Info("收到客户端的SettingFrame")
	hasWebTransport := false
	for key, val := range settingFrameFromClient.Payload {
		if h3.SettingEnableWebTransport == key && 1 == val {
			hasWebTransport = true
		}
	}
	if !hasWebTransport {
		log.Error("客户端不支持webtransport,关闭")
		return
	}

	// 等待客户端的bidStream
	bidStreamFromClient, acceptBidStreamFromClientErr := s.quicSession.AcceptStream(ctx)
	if nil != acceptBidStreamFromClientErr {
		log.Error("accept BidStreamFromClient error", zap.Error(acceptBidStreamFromClientErr))
		return
	}
	log.Info("收到客户端的BidStream")
	// 读取headers
	headersFrame := h3.ParseHeadersFrame(bidStreamFromClient)
	log.Info("收到客户端的headers", zap.String("Protocol", headersFrame.Protocol),
		zap.String("Authority", headersFrame.Authority),
		zap.String("Path", headersFrame.Path))
	headersFrameFromServer := h3.HeadersFrame{Header: map[string][]string{
		"sec-webtransport-http3-draft": {"draft02"},
	}}
	_, _ = bidStreamFromClient.Write(headersFrameFromServer.ToBytes(200))
	log.Info("服务端headersFrame发送结束")

	// 等待客户端的bidStream
	s.HandleBidStream(headersFrame.Header.Get("Origin"), headersFrame.Path)
}

func (s *Session) HandleBidStream(origin, path string) {
	ctx := s.quicSession.Context()
	log := s.log
	bidStream, acceptBidStreamErr := s.quicSession.AcceptStream(ctx)
	if nil != acceptBidStreamErr {
		log.Error("custom accept BidStream error", zap.Error(acceptBidStreamErr))
	} else {
		s.bidStream = bidStream
		// 读取Type
		reader := quicvarint.NewReader(bidStream)
		sType, readTypeErr := quicvarint.Read(reader)
		if nil != readTypeErr || 0x41 != sType {
			log.Error("custom bidStream type 异常", zap.Error(readTypeErr))
			return
		}
		// 读取SessionId
		_, readSessionIdErr := quicvarint.Read(reader)
		if nil != readSessionIdErr {
			log.Error("read bidStream sessionId error", zap.Error(readSessionIdErr))
			return
		}

		s.signal.fireConnectCb(s)

		if !s.signal.fireAuthCb(proto.AuthCbStruct{
			Path:   path,
			Origin: origin,
		}) {
			return
		}

		// 读取数据
		lenHeaderBytes := make([]byte, 4)
		lenBodyBytes := make([]byte, 4)
		for {
			// 读取 header 长度
			_, readDataErr := io.ReadFull(reader, lenHeaderBytes)
			if nil != readDataErr {
				log.Error("read bidStream data error", zap.Error(readDataErr))
				break
			}
			// 读取 body 长度
			_, readDataErr = io.ReadFull(reader, lenBodyBytes)
			if nil != readDataErr {
				log.Error("read bidStream data error", zap.Error(readDataErr))
				break
			}

			lenHeader := proto.BytesToInt(lenHeaderBytes)
			lenBody := proto.BytesToInt(lenBodyBytes)

			leftBuf := make([]byte, lenHeader+lenBody)
			_, readDataErr = io.ReadFull(reader, leftBuf)
			if nil != readDataErr {
				log.Error("read bidStream data error", zap.Error(readDataErr))
				break
			}
			totalBuf := bytes.Buffer{}
			totalBuf.Write(lenHeaderBytes)
			totalBuf.Write(lenBodyBytes)
			totalBuf.Write(leftBuf)

			// 解析数据包
			structFromBytes := proto.NewMsgStructFromBytes(totalBuf.Bytes())
			if nil != structFromBytes {
				s.fireOnCb(structFromBytes)
			}
		}

	}
}

func (s *Session) fireOnCb(msg *proto.MsgStruct) {
	header := msg.GetHeader()
	cbsRaw, ok := s.onCbMap.Load(header)
	if ok {
		for _, cb := range cbsRaw.([]proto.SessionCallback) {
			go cb(proto.SessionCbStruct{
				Session:   s,
				MsgStruct: msg,
				SessionId: s.id,
				Msg:       msg.GetBody(),
			})
		}
	}
	cbsRaw, ok = s.onCbMap.Load("*")
	if ok {
		for _, cb := range cbsRaw.([]proto.SessionCallback) {
			go cb(proto.SessionCbStruct{
				Session:   s,
				MsgStruct: msg,
				SessionId: s.id,
				Msg:       msg.GetBody(),
			})
		}
	}
}
